#include "fft2d.h"

#include <math.h>
#include <string.h>

using namespace fft2d;

Complex_array::Complex_array(int nx, int ny)
:
nx(nx), ny(ny)
{
	data = new Complex[nx*ny];
	for (int i = 0 ; i < nx*ny ; ++i)
	{
		data[i]= 0;
	}
}

Complex_array::~Complex_array()
{
	delete [] data;
	data = 0;
}

Complex_array::Complex_array(const Complex_array & orig)
:
nx(orig.nx), ny(orig.ny)
{
	data = new Complex[nx*ny];
	memcpy(data, orig.data, nx*ny*sizeof(Complex));
}

Complex_array & Complex_array::operator=(const Complex_array & rhs)
{
	delete [] data;
	nx = rhs.nx;
	ny = rhs.ny;
	data = new Complex[nx*ny];
	memcpy(data, rhs.data, nx*ny*sizeof(Complex));
	return *this;
}

void Complex_array::rotate()
{
	Complex_array orig(*this);
	int nx2 = nx/2;
	int ny2 = ny/2;
	assert(nx2*2 == nx);
	assert(ny2*2 == ny);
	for (int i = 0 ; i < nx ; ++i)
	{
		for (int j = 0 ; j < ny ; ++j)
		{
			(*this)(i, j) = orig( (i+nx2) % nx, (j+ny2) % ny );
		}
	}
}

void Complex_array::flip_lr()
{
	Complex_array orig(*this);
	for (int i = 0 ; i < nx ; ++i)
	{
		for (int j = 0 ; j < ny ; ++j)
		{
			(*this)(i, j) = orig( (nx-1) - i, j );
		}
	}
}

void Complex_array::flip_ud()
{
	Complex_array orig(*this);
	for (int i = 0 ; i < nx ; ++i)
	{
		for (int j = 0 ; j < ny ; ++j)
		{
			(*this)(i, j) = orig( i, (ny -1) - j );
		}
	}
}


void Complex_array::tile()
{
	Complex_array bot_left(*this); // stays in same location
	Complex_array bot_right(*this);    
	bot_right.flip_lr();
	Complex_array top_left(*this);
	top_left.flip_ud();
	Complex_array top_right(top_left);
	top_right.flip_lr();

	delete [] data;
	int nx0 = nx;
	int ny0 = ny;
	nx *= 2;
	ny *= 2;

	data = new Complex[nx*ny];

	int i, j;
	for (i = 0 ; i < nx0 ; ++i)
	{
		for (j = 0 ; j < ny0 ; ++j)
		{
			(*this)(i,     j    ) = bot_left(i, j);
			(*this)(i+nx0, j    ) = bot_right(i, j);
			(*this)(i,     j+ny0) = top_left(i, j);
			(*this)(i+nx0, j+ny0) = top_right(i, j);
		}
	}
}


#if defined(__APPLE__) || defined(MACOSX) || defined(unix)
#else
int rint(double val)
{
	if (val > 0)
		return (int) (val + 0.5);
	else
		return (int) (val - 0.5);
}
#endif

/* 
returns true if nx is a power of two, and sets m so that 2^m = nx 
*/
bool powerof2(int nx, int * m)
{
	*m = (int) rint( log((double) nx)/log(2.0));
	if (rint(pow(2.0, (double) *m)) != nx)
	{
		return false;
	}
	return true;
}

/*-------------------------------------------------------------------------
This computes an in-place complex-to-complex FFT
x and y are the real and imaginary arrays of 2^m points.
dir =  1 gives forward transform	  0
dir = -1 gives reverse transform

Formula: forward
N-1
---
1   \          - j k 2 pi n / N
X(n) = ---   > x(k) e                    = forward transform
N   /                                n=0..N-1
---
k=0

Formula: reverse
N-1
---
\          j k 2 pi n / N
X(n) =       > x(k) e                    = forward transform
/                                n=0..N-1
---
k=0
*/
bool FFT(int dir,int m,float *x,float *y)
{
	long nn,i,i1,j,k,i2,l,l1,l2;
	float c1,c2,tx,ty,t1,t2,u1,u2,z;

	/* Calculate the number of points */
	nn = 1;
	for (i=0;i<m;i++)
		nn *= 2;

	/* Do the bit reversal */
	i2 = nn >> 1;
	j = 0;
	for (i=0;i<nn-1;i++) 
	{
		if (i < j) 
		{
			tx = x[i];
			ty = y[i];
			x[i] = x[j];
			y[i] = y[j];
			x[j] = tx;
			y[j] = ty;
		}
		k = i2;
		while (k <= j) 
		{
			j -= k;
			k >>= 1;
		}
		j += k;
	}

	/* Compute the FFT */
	c1 = -1.0;
	c2 = 0.0;
	l2 = 1;
	for (l=0;l<m;l++) 
	{
		l1 = l2;
		l2 <<= 1;
		u1 = 1.0;
		u2 = 0.0;
		for (j=0;j<l1;j++) 
		{
			for (i=j;i<nn;i+=l2) 
			{
				i1 = i + l1;
				t1 = u1 * x[i1] - u2 * y[i1];
				t2 = u1 * y[i1] + u2 * x[i1];
				x[i1] = x[i] - t1;
				y[i1] = y[i] - t2;
				x[i] += t1;
				y[i] += t2;
			}
			z =  u1 * c1 - u2 * c2;
			u2 = u1 * c2 + u2 * c1;
			u1 = z;
		}
		c2 = (float) sqrt((1.0 - c1) / 2.0);
		if (dir == 1)
			c2 = -c2;
		c1 = (float) sqrt((1.0 + c1) / 2.0);
	}

	/* Scaling for forward transform */
	if (dir == 1) 
	{
		for (i=0;i<nn;i++) 
		{
			x[i] /= nn;
			y[i] /= nn;
		}
	}

	return true;
}

/*-------------------------------------------------------------------------
Perform a 2D FFT inplace given a complex 2D array
The direction dir, 1 for forward, -1 for reverse
The size of the array (nx,ny)
Return false if there are memory problems or
the dimensions are not powers of 2
*/
bool fft2d::FFT2D(Complex_array & c,int nx,int ny,int dir)
{
	int i,j;
	int m;
	float *real,*imag;

	/* Transform the rows */
	real = new float[nx];
	imag = new float[nx];

	if (!powerof2(nx,&m))
	{
		delete [] real;
		delete [] imag;
		return false;
	}

	for (j=0;j<ny;j++) 
	{
		for (i=0;i<nx;i++) 
		{
			real[i] = c(i, j).real;
			imag[i] = c(i, j).imag;
		}
		FFT(dir,m,real,imag);
		for (i=0;i<nx;i++) 
		{
			c(i, j).real = real[i];
			c(i, j).imag = imag[i];
		}
	}

	delete [] real;
	delete [] imag;

	/* Transform the columns */
	real = new float[ny];
	imag = new float[ny];

	if (!powerof2(ny,&m))
	{
		delete [] real;
		delete [] imag;
		return false;
	}

	for (i=0;i<nx;i++) 
	{
		for (j=0;j<ny;j++) 
		{
			real[j] = c(i, j).real;
			imag[j] = c(i, j).imag;
		}
		FFT(dir,m,real,imag);
		for (j=0;j<ny;j++) 
		{
			c(i, j).real = real[j];
			c(i, j).imag = imag[j];
		}
	}
	delete [] real;
	delete [] imag;

	return true;
}
